from 读取文件 import *
import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import solve


def CF(x, y, n=6):  # 输入要拟合的两个向量, 得到向量形式的多项式
    x = np.array(x)
    y = np.mat(y).T
    C_matrix = np.mat(np.ones(len(x))).T  # 初始化C_matrix作为系数矩阵,
    k = 1
    while k <= n:
        a = np.mat(x ** k).T  # 临时变量abc
        C_matrix = np.append(a, C_matrix, axis=1)
        k += 1
        pass  # 得到的C_matrix就是系数矩阵
    # 接下来求这个不一致方程组的最小二乘解
    A = np.dot(C_matrix.T, C_matrix)
    B = np.dot(C_matrix.T, y)
    x_ = np.linalg.solve(A, B)
    Error = y - C_matrix * x_

    return {"Factor": x_, "Error": Error}


if __name__ == "__main__":  # 测试代码
    fig = plt.figure(figsize=(6, 6))

    t = np.linspace(0, 1.06, 1000)
    r = 50
    w = 1
    x = r * (np.cos(w * t + 1.082) + w * t * np.sin(w * t + 1.082))
    x = x - x[0]
    y = r * (np.sin(w * t + 1.082) - w * t * np.cos(w * t + 1.082))
    y = y - y[0]

    re1 = CF(y, x)
    s = np.linspace(0, 30, 1000)
    plt.scatter(x, y, color='blue')
    plt.plot(np.polyval(re1["Factor"], s), s, color='red')
    print(re1["Error"])
    # plt.plot(re1[1])
    plt.show()
    pass
